from langchain_postgres.vectorstores import PGVector
from langchain_ollama import OllamaEmbeddings, OllamaLLM
import os, time, json

# See docker command above to launch a postgres instance with pgvector enabled.
connection = "postgresql+psycopg://langchain:langchain@192.168.137.2:5432/langchain"  # Uses psycopg3!
collection_name = "my_docs"

OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'http://127.0.0.1:11434')

vector_store = PGVector(
    embeddings= OllamaEmbeddings(
            model='mxbai-embed-large',
            base_url=OLLAMA_HOST  # 添加base_url参数
        ),
    collection_name=collection_name,
    connection=connection,
    use_jsonb=True,
)

# # 增加文档
# from langchain_core.documents import Document

# document_1 = Document(page_content="foo", metadata={"baz": "bar"})
# document_2 = Document(page_content="thud", metadata={"bar": "baz"})
# document_3 = Document(page_content="i will be deleted :(")

# documents = [document_1, document_2, document_3]
# ids = ["1", "2", "3"]
# vector_store.add_documents(documents=documents, ids=ids)

# # 删除
# vector_store.delete(ids=["3"])

# 查询
# results = vector_store.similarity_search(query="thud",k=1)
# print(results)
# for doc in results:
#     print(f"* {doc.page_content} [{doc.metadata}]")

# 过滤
# results = vector_store.similarity_search(query="thud",k=1,filter={"bar": "baz"})
# for doc in results:
#     print(f"* {doc.page_content} [{doc.metadata}]")

# results = vector_store.similarity_search_with_score(query="qux",k=1)
# for doc, score in results:
#     print(f"* [SIM={score:3f}] {doc.page_content} [{doc.metadata}]")

# 
retriever = vector_store.as_retriever(
    search_type="mmr",
    search_kwargs={"k": 1, "fetch_k": 2, "lambda_mult": 0.5},
)
results = retriever.invoke("thud")
print(results)
for doc in results:
    print(f"* {doc.page_content} [{doc.metadata}]")

